import os
import shutil
import json

path_04 = '/tcdata/tile_round2_train_20210204_update/train_imgs'
path_08 = '/tcdata/tile_round2_train_20210208/train_imgs'

temp_04 = '/tcdata/tile_round2_train_20210204_update/train_template_imgs'
temp_08 = '/tcdata/tile_round2_train_20210208/train_template_imgs'


target_path = '/my_data'
if not os.path.exists(target_path):
    os.mkdir(target_path)

def data_combine():
    if not os.path.exists(os.path.join(target_path,'ori_img' )):
        os.mkdir(os.path.join(target_path,'ori_img' ))
    if not os.path.exists(os.path.join(target_path,'ori_temp' )):
        os.mkdir(os.path.join(target_path,'ori_temp' ))


    for file in os.listdir(path_04):
        sour = os.path.join(path_04, file)
        targ = os.path.join(target_path,'ori_img',file)
        shutil.copy(sour, targ)

    for file in os.listdir(path_08):
        sour = os.path.join(path_08, file)
        targ = os.path.join(target_path,'ori_img',file)
        shutil.copy(sour, targ)


    for file in os.listdir(temp_04):
        sour = os.path.join(temp_04, file)
        targ = os.path.join(target_path,'ori_temp',file)
        shutil.copy(sour, targ)

    for file in os.listdir(temp_08):
        sour = os.path.join(temp_08, file)
        targ = os.path.join(target_path,'ori_temp',file)
        shutil.copy(sour, targ)



def anno_combine():
    anno_04 = json.load(open('/tcdata/tile_round2_train_20210204_update/train_annos.json', 'r'))
    anno_08 = json.load(open('/tcdata/tile_round2_train_20210208/train_annos.json', 'r'))

    anno_04.extend(anno_08)
    all_anno = anno_04
    all_anno = json.dumps(all_anno)
    with open(os.path.join(target_path, 'train_annos.json'), 'w') as f:
        f.write(all_anno)

def ann2coco():
    ori_anno = json.load(open(os.path.join(target_path,'train_annos.json' ),'r'))
    exist_imgs = []
    imgs = []
    annotations = []
    categories = [{'id':1,'name':"edge"},{'id':2,'name':'corner'},{'id':3,'name':'white'},{'id':4,'name':'light'},{'id':5,'name':'deep'},{'id':6,'name':'circle'}, {'id':7, 'name':'mark'}, {'id': 8, 'name': 'scrach'}]
    num_img = 0
    for dictt in ori_anno:
        file_name = dictt['name']
        if file_name not in exist_imgs:
            num_img += 1
            exist_imgs.append(file_name)
            imgs.append({'file_name':file_name,'height':dictt['image_height'],'width':dictt['image_width'],'id':num_img})
    ii=0
    for dictt in ori_anno:
        ii += 1
        file_name = dictt['name']
        for img in imgs:
            if img['file_name'] == file_name:
                iid = img['id']
                break
        box = dictt['bbox']
        annotations.append({'image_id':iid,'bbox':[box[0],  box[1],   box[2]-box[0],     box[3]-box[1] ],'area': (box[2]-box[0])*( box[3]-box[1]) ,'category_id':dictt['category'],'id':ii})

    all = {'images':imgs,'annotations':annotations,'categories':categories}

    all = json.dumps(all)
    with open(os.path.join(target_path, 'annotations_cocoFormat.json'),'w') as f:
        f.write(all)


if __name__ == '__main__':
    data_combine()
    anno_combine()
    ann2coco()